package tt.langchain4j.service;

import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import tt.langchain4j.core.BufferSupplier;
import tt.langchain4j.model.ChatLanguageModels;
import tt.langchain4j.model.EmbeddingModels;
import tt.langchain4j.rag.MultiableRetribevalAugmentor;
import tt.langchain4j.store.EmbeddingStores;

import java.util.function.Supplier;

public enum ChatServices {


    DEEP_SEEK_R1_7B(()->AiChatServiceFactory.builder()
            .chatLanguageModel(ChatLanguageModels.DEEP_SEEK_R1_7B.model())
            .streamingChatLanguageModel(ChatLanguageModels.DEEP_SEEK_R1_7B.streamModel())
            .build()

    ) ,

    DEEP_SEEK_R1_7B_WITH_RAG(()->AiChatServiceFactory.builder()
            .chatLanguageModel(ChatLanguageModels.DEEP_SEEK_R1_7B.model())
            .streamingChatLanguageModel(ChatLanguageModels.DEEP_SEEK_R1_7B.streamModel())
            .retrievalAugmentor(MultiableRetribevalAugmentor.from(
                    DefaultRetrievalAugmentor.builder()
                            .contentRetriever(EmbeddingStoreContentRetriever
                                    .builder()
                                    .embeddingModel(EmbeddingModels.DEFAULT.model())
                                    .embeddingStore(EmbeddingStores.DEFAULT.embeddingStore())
                                    .minScore(0.8)
                                    .build())
                            .build()

            ))
            .build()

    ) ,

    LLAVA(()->AiChatServiceFactory.builder()
            .chatLanguageModel(ChatLanguageModels.LLAVA.model())
            .streamingChatLanguageModel(ChatLanguageModels.LLAVA.streamModel())
            .build()
    ) ,

    LLAMA_32_VISION_11B(()-> AiChatServiceFactory.builder()
            .chatLanguageModel(ChatLanguageModels.LLAMA_32_VISION_11B.model())
            .streamingChatLanguageModel(ChatLanguageModels.LLAMA_32_VISION_11B.streamModel())
            .build()
    ) ,


    ;

    private final BufferSupplier<AiChatService> serviceSupplier;

    ChatServices(Supplier<AiChatService> serviceSupplier) {
        this.serviceSupplier = new BufferSupplier<>(serviceSupplier);
    }


    public AiChatService chatService() {
        return serviceSupplier.get();
    }
}
